{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tuning Llama"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction\n",
    "\n",
    "In this tutorial, we'll focus on how you can easily tune various flavors of Llama. For simplicity, we'll be using Polaris as the platform for compute.\n",
    "\n",
    "**NOTE:** This tutorial builds off of our [Finetuning Tutorial](https://github.com/oumi-ai/oumi/blob/main/notebooks/Oumi%20-%20Finetuning%20Tutorial.ipynb). We recommend starting there first to get a thorough understanding of how tuning works in our library."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prerequisites"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tutorial assumes:\n",
    "- You have a valid ALCF account with access to Polaris\n",
    "- You're familiar with our tuning flow\n",
    "- You're familiar with how to launch Oumi workflows on Polaris. [Here's a relevant tutorial](https://github.com/oumi-ai/oumi/blob/main/notebooks/Oumi%20-%20Deploying%20a%20Job.ipynb)\n",
    "- You've signed [Llama's agreement on HuggingFace](https://huggingface.co/meta-llama/Llama-3.1-8B)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tuning\n",
    "\n",
    "We currently have out-of-the-box Polaris tuning job configs for the following flavors of Llama:\n",
    "\n",
    "- Llama3.1 8B LoRA: [configs/recipes/llama3_1/sft/8b_lora/polaris_job.yaml](https://github.com/oumi-ai/oumi/blob/main/configs/recipes/llama3_1/sft/8b_lora/polaris_job.yaml) ✨\n",
    "  - Tuning config: [configs/recipes/llama3_1/sft/8b_lora/train.yaml](https://github.com/oumi-ai/oumi/blob/main/configs/recipes/llama3_1/sft/8b_lora/train.yaml)\n",
    "- Llama3.1 8B SFT: [configs/recipes/llama3_1/sft/8b_full/polaris_job.yaml](https://github.com/oumi-ai/oumi/blob/main/configs/recipes/llama3_1/sft/8b_full/polaris_job.yaml) ✨\n",
    "  - Tuning config: [configs/recipes/llama3_1/sft/8b_full/train.yaml](https://github.com/oumi-ai/oumi/blob/main/configs/recipes/llama3_1/sft/8b_full/train.yaml)\n",
    "- Llama3.1 70B LoRA: [configs/recipes/llama3_1/sft/70b_lora/polaris_job.yaml](https://github.com/oumi-ai/oumi/blob/main/configs/recipes/llama3_1/sft/70b_lora/polaris_job.yaml) ✨\n",
    "  - Tuning config: [configs/recipes/llama3_1/sft/70b_lora/train.yaml](https://github.com/oumi-ai/oumi/blob/main/configs/recipes/llama3_1/sft/70b_lora/train.yaml)\n",
    "- Llama3.1 70B SFT: [configs/recipes/llama3_1/sft/70b_full/polaris_job.yaml](https://github.com/oumi-ai/oumi/blob/main/configs/recipes/llama3_1/sft/70b_full/polaris_job.yaml) ✨\n",
    "  - Tuning config: [configs/recipes/llama3_1/sft/70b_full/train.yaml](https://github.com/oumi-ai/oumi/blob/main/configs/recipes/llama3_1/sft/70b_full/train.yaml)\n",
    "\n",
    "By default our tuning job will run using the [`yahma/alpaca-cleaned`](https://huggingface.co/datasets/yahma/alpaca-cleaned) dataset. This is configured in the tuning configs above. We strongly suggest tuning these parameters as needed for your specific run.\n",
    "\n",
    "Before running the job, ensure you've signed Llama's agreement on HuggingFace and have obtained your [HF_TOKEN](https://huggingface.co/docs/huggingface_hub/en/package_reference/environment_variables#hftoken). We'll pass it to our job by appending `envs.HF_TOKEN=$HF_TOKEN` to our launcher script. This is not currently needed as our script copies the model from Polaris' Eagle file system, but it will be needed to pull the model directly from HuggingFace.\n",
    "\n",
    "The example command below uses the `preemptable` queue, which supports up to 10 nodes and up to 3 days runtime. When debugging, prefer using the `debug` (for <=2 node jobs) or `debug-scaling` (for <=10 node jobs) queues, which usually has less queue time. However, you need to modify the `walltime` in the yaml config above to be <= 1 hour. For more details on available queues, see [ALCF's documentation](https://docs.alcf.anl.gov/polaris/running-jobs/)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "vscode": {
     "languageId": "raw"
    }
   },
   "source": [
    "```shell\n",
    "# Replace with your desired config. We use Llama 8B LoRA below as it's the fastest.\n",
    "OUMI_CONFIG_PATH=\"configs/recipes/llama3_1/sft/8b_lora/polaris_job.yaml\"\n",
    "# If using debug/debug-scaling queue, make sure that the walltime is <= 1 hr.\n",
    "POLARIS_QUEUE=\"preemptable\"\n",
    "oumi launch up -c $OUMI_CONFIG_PATH --cluster $POLARIS_QUEUE.$ALCF_USER --envs.HF_TOKEN $HF_TOKEN --user $ALCF_USER\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluation\n",
    "\n",
    "Now that you've tuned your model, you can evaluate your tuned model using EleutherAI's Language Model Eval Harness (LM Harness). We currently have the following eval configs which by default evaluate MMLU English:\n",
    "\n",
    "- Llama3.1 8B: [configs/recipes/llama3_1/evaluation/8b_polaris_job.yaml](https://github.com/oumi-ai/oumi/blob/main/configs/recipes/llama3_1/evaluation/8b_polaris_job.yaml) ✨\n",
    "- Llama3.1 70B: [configs/recipes/llama3_1/evaluation/70b_polaris_job.yaml](https://github.com/oumi-ai/oumi/blob/main/configs/recipes/llama3_1/evaluation/70b_polaris_job.yaml) ✨\n",
    "\n",
    "Make sure to modify `MODEL_CHECKPOINT_DIR` in the eval config to your desired checkpoint, for either a full fine-tuned or LoRA-tuned model. The above tuning run should output to a directory in `/eagle/community_ai/<your_alcf_username>/runs`. For example, for Llama 8B LoRA tuning, it would output to `/eagle/community_ai/<your_alcf_username>/runs/llama8b.lora.<pbs_job_id>`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```shell\n",
    "# Replace with your desired config. We use Llama 8B eval below.\n",
    "OUMI_EVAL_CONFIG_PATH=\"configs/recipes/llama3_1/evaluation/8b_polaris_job.yaml\"\n",
    "# If using debug/debug-scaling queue, make sure that the walltime is <= 1 hr.\n",
    "POLARIS_QUEUE=\"debug\"\n",
    "oumi launch up -c $OUMI_EVAL_CONFIG_PATH --cluster $POLARIS_QUEUE.$ALCF_USER --envs.HF_TOKEN $HF_TOKEN --user $ALCF_USER\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
